from fetchingthread import FetchingThread
#from time import sleep
from iphelper import StrToInt
from certparser import CertParser
from M2Crypto import X509
from sqlstore import DB
from logger import Logger
from reporter import Reporter

Log = Logger().Log
R = Reporter()

class ThreadManager():
    def __init__(self, threads=0): # threads=0 -> estimate based on ip range
        "This is main thread. All tasks except network connection should be done here."
        self.threads = threads
        self.workers = []
        self.cert_buffer = []
        self.DB = DB()
        self.worker_stat_table = R.WorkerTable()
        self.global_stats = R.GlobalStatsTable()
    def __repr__(self):
        return "ThreadManager class"
    def SetThreadsNum(self, threads):
        Log("Setting number of network threads to: %d" % threads, 4)
        self.threads = threads
    def EstimateThreadsNum(self, start_ip, stop_ip):
        up_limit = 256 # move this to configuration file
        down_limit = 1 # changing this may cause problems with small num of ips to scan.
        threads_num = (stop_ip - start_ip) / 8 # todo: improve this algorithm
        if threads_num > up_limit:
            threads_num = up_limit
        elif threads_num < down_limit:
            threads_num = down_limit
        return threads_num
    def MakeWorkersTable(self, start_ip, stop_ip):
        for i in range(self.threads):
            self.worker_stat_table.append({'hits' : 0, 'fetched' : 0, 'failed' : 0})
            self.workers.append(FetchingThread(start_ip, stop_ip, i, self.threads, self.cert_buffer, self.worker_stat_table))
    def AddJobStats(self, start_ip, stop_ip):
        self.global_stats['targets_num'] = stop_ip - start_ip
        self.DB.insert_job(start_ip, stop_ip)
    def UpdateJobStats(self):
        self.DB.end_job(self.global_stats['fetched'])
    def PurgeCert(self, cert):
        self.cert_buffer.remove(cert)
    def WaitOrQuit(self, timeout):
        for worker in self.workers:
            if worker.is_alive():
                worker.join(timeout)
                return
        Log(R.MakeSummary())
        self.UpdateJobStats()
        Log('Goodbye.')
        exit(0)
    def ParseAndInsert(self, cert):
        ""
        CP = CertParser()
        target = list(cert[:2])  #host, port
        x509_pem = cert[2] #x509 PEM formated
        x509 = X509.load_cert_string(x509_pem, X509.FORMAT_PEM)
        rsa = CP.ExtractPubKey(x509)
        x509_attr = CP.ParseX509(x509)
        try:
            rsa_attr = CP.ParsePubKey(rsa)
            all = target + [x509_pem] + x509_attr + rsa_attr
            self.DB.insert_x509(all)
        except ValueError: #ValueError("PKey instance is not holding a RSA key - probably DSA")
            #rsa_attr = None #Unneeded
            Log('Exception caught during RSA extraction. X509 PEM: %s:%s\n%s' % (target[0], target[1], x509_pem), 5)
    def Start(self, start_ip, stop_ip):
        ""
        start_ip, stop_ip = StrToInt(start_ip), StrToInt(stop_ip) #convert dot separated ip into integer
        if not self.threads:
            threads = self.EstimateThreadsNum(start_ip, stop_ip)
            self.SetThreadsNum(threads)
        if not self.workers:
            self.MakeWorkersTable(start_ip, stop_ip)
        self.AddJobStats(start_ip, stop_ip)
        for i in range(len(self.workers)):
            self.workers[i].start()
        while (1):
            for cert in self.cert_buffer:
                Log("X509 certificate obtained from: %s:%s" % (cert[0], cert[1]), 4)
                self.ParseAndInsert(cert)
                self.PurgeCert(cert)
            self.WaitOrQuit(2)
